// Copyright 2021 Twitter, Inc.
// Licensed under the Apache License, Version 2.0
// http://www.apache.org/licenses/LICENSE-2.0

use async_trait::async_trait;

use std::path::Path;
use std::sync::{Arc, Mutex};

use crate::common::bpf::BPF;
use crate::config::SamplerConfig;
use crate::samplers::{Common, Sampler};
use crate::*;

#[cfg(feature = "bpf")]
use crate::common::bpf::bpf_hash_char_to_map;

mod config;
mod stat;

pub use config::{LibraryProbeConfig, UsercallConfig, NAMESPACE};
pub use stat::UsercallStatistic;

#[allow(dead_code)]
pub struct Usercall {
    bpf: Option<Arc<Mutex<BPF>>>,
    bpf_last: Arc<Mutex<Instant>>,
    common: Common,
    statistics: Vec<UsercallStatistic>,
    libraries: Vec<LibraryProbeConfig>,
}

const DEFAULT_LIB_SEARCH_PATHS: [&str; 6] = [
    "/lib64",
    "/usr/lib64",
    "/usr/local/lib64",
    "/lib",
    "/usr/lib",
    "/usr/local/lib",
];

#[cfg(feature = "bpf")]
const PROBE_PRELUDE: &str = r#"
#include <uapi/linux/ptrace.h>
struct key_t {
    char c[80];
};
BPF_HASH(counts, struct key_t);

"#;

macro_rules! probe_template {
    () => {
        r#"

int probe_{}(void *ctx) {{
    struct key_t key = {{.c = "{}"}};
    u64 zero = 0, *val;
    val = counts.lookup_or_init(&key, &zero);
    (*val)++;
    return 0;
}}
"#
    };
}

/// Determines if a path is a match for the given lib_name
///
/// This custom logic is necessary because of the common convention of versioning linux libraries by
/// appending the version number to the library file name. I.E. the file /usr/lib/libkrb5.so.3.3
/// results in an extension of '3' with a naive parsing of the file extension. This is what Path
/// and PathBuf would return for the extension.
#[allow(dead_code)]
fn path_match(lib_name: &str, path: &Path) -> bool {
    if let Some(file_name) = path.file_name() {
        if let Some(file_str) = file_name.to_str() {
            let parts: Vec<&str> = file_str.split('.').collect();
            if parts.len() < 2 {
                return false;
            }
            let mut stem_str: String = parts[0].to_string();
            let mut ext_str: String = parts[1].to_string();
            let end_index = parts.len() - 1;
            if parts[end_index] == "so" {
                stem_str = parts[..end_index].join(".");
                ext_str = "so".into();
            }
            let to_test = match stem_str.starts_with("lib") {
                true => stem_str[3..].into(),
                false => stem_str,
            };
            return to_test.eq(lib_name) && "so".eq(&ext_str[..]);
        }
    }
    false
}

impl Usercall {
    fn init_bpf(&mut self) -> Result<(), anyhow::Error> {
        // The bpf source code generated by probe_template
        let mut bpf_probes = String::new();

        // The list of probes that have been found.
        let mut found_probes: Vec<(String, String, String)> = Vec::new();

        // Add probes that are linked to specific files in the config
        for probe_config in self.libraries.iter().filter(|x| x.path.is_some()) {
            for func in probe_config.functions.iter() {
                #[allow(clippy::format_in_format_args)]
                bpf_probes.push_str(&format!(
                    probe_template!(),
                    found_probes.len(),
                    format!("{}/{}/{}", NAMESPACE, probe_config.name, func)
                ));
                found_probes.push((
                    probe_config.path.as_ref().unwrap().clone(),
                    probe_config.name.clone(),
                    func.clone(),
                ));
            }
        }

        // Add probes by searching the default paths.
        let entries: Vec<walkdir::DirEntry> = DEFAULT_LIB_SEARCH_PATHS
            .iter()
            .map(|p| {
                walkdir::WalkDir::new(p)
                    .follow_links(true)
                    .into_iter()
                    .filter_map(|e| e.ok())
            })
            .flatten()
            .collect();

        for probe_config in self.libraries.iter().filter(|x| x.path.is_none()) {
            let mut found = false;
            for entry in &entries {
                if path_match(&probe_config.name, entry.path()) {
                    for func in probe_config.functions.iter() {
                        #[allow(clippy::format_in_format_args)]
                        bpf_probes.push_str(&format!(
                            probe_template!(),
                            found_probes.len(),
                            format!("{}/{}/{}", NAMESPACE, probe_config.name, func)
                        ));
                        found_probes.push((
                            entry.path().to_string_lossy().to_string(),
                            probe_config.name.clone(),
                            func.clone(),
                        ));
                    }
                    found = true;
                    break;
                }
            }
            if !found {
                let err = std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!("probe not found {:?}", probe_config),
                );
                if self.common.config().fault_tolerant() {
                    for func in probe_config.functions.iter() {
                        let stat_name = format!("{}/{}/{}", NAMESPACE, probe_config.name, func);
                        self.statistics.retain(|s| s.stat_path != stat_name);
                    }
                    warn!("{}", err);
                } else {
                    return Err(err.into());
                }
            }
        }

        #[cfg(feature = "bpf")]
        {
            debug!("Registering probes: {:?}", found_probes);
            // Build the bpf program by appending all the bpf_probe source to the prelude
            let bpf_prog = PROBE_PRELUDE.to_string() + &bpf_probes;
            let mut bpf = bcc::BPF::new(&bpf_prog)?;
            for (i, probe) in found_probes.iter().enumerate() {
                let (path, lib, func) = probe;
                if let Err(e) = bcc::Uprobe::new()
                    .handler(&format!("probe_{}", i))
                    .binary(path)
                    .symbol(func)
                    .attach(&mut bpf)
                {
                    let err = match std::fs::metadata(path) {
                        Ok(md) => {
                            if md.is_file() {
                                std::io::Error::new(
                                    std::io::ErrorKind::Other,
                                    format!("problem probing {:?}: {}", probe, e),
                                )
                            } else {
                                std::io::Error::new(
                                    std::io::ErrorKind::NotFound,
                                    format!("problem probing {:?}: Not a file: {}", probe, path),
                                )
                            }
                        }
                        Err(e) => std::io::Error::new(
                            std::io::ErrorKind::Other,
                            format!("problem probing {:?}: {}", probe, e),
                        ),
                    };
                    if self.common.config().fault_tolerant() {
                        let stat_name = format!("{}/{}/{}", NAMESPACE, lib, func);
                        self.statistics.retain(|s| s.stat_path != stat_name);
                        warn!("{}", err);
                    } else {
                        Err(err)?
                    }
                };
            }

            self.bpf = Some(Arc::new(Mutex::new(BPF { inner: bpf })));
        }

        Ok(())
    }
}

#[async_trait]
impl Sampler for Usercall {
    type Statistic = UsercallStatistic;

    fn new(common: Common) -> Result<Self, anyhow::Error> {
        let statistics = common.config().samplers().usercall().statistics();
        let libraries = common.config().samplers().usercall().libraries();

        let mut sampler = Self {
            bpf: None,
            bpf_last: Arc::new(Mutex::new(Instant::now())),
            common,
            statistics,
            libraries,
        };
        if sampler.sampler_config().enabled() {
            sampler.init_bpf()?;
            sampler.register();
        }
        Ok(sampler)
    }

    fn spawn(common: Common) {
        if common.config().samplers().usercall().enabled() {
            match Self::new(common.clone()) {
                Ok(mut sampler) => {
                    common.runtime().spawn(async move {
                        loop {
                            let _ = sampler.sample().await;
                        }
                    });
                }
                Err(e) => {
                    if !common.config.fault_tolerant() {
                        fatal!("failed to initialize usercall sampler {}", e);
                    } else {
                        error!("failed to initialize usercall sampler {}", e);
                    }
                }
            }
        }
    }

    fn common(&self) -> &Common {
        &self.common
    }

    fn common_mut(&mut self) -> &mut Common {
        &mut self.common
    }

    fn sampler_config(&self) -> &dyn SamplerConfig<Statistic = Self::Statistic> {
        self.common.config().samplers().usercall()
    }

    async fn sample(&mut self) -> Result<(), std::io::Error> {
        if let Some(ref mut delay) = self.delay() {
            delay.tick().await;
        }

        if !self.sampler_config().enabled() {
            return Ok(());
        }

        #[cfg(feature = "bpf")]
        if let Some(ref bpf) = self.bpf {
            let bpf = bpf.lock().unwrap();
            let table = (*bpf)
                .inner
                .table("counts")
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
            let stat_map = bpf_hash_char_to_map(&table);
            for stat in self.statistics.iter() {
                let val = stat_map.get(&stat.stat_path).unwrap_or(&0);
                self.metrics()
                    .record_counter(stat, Instant::now(), *val)
                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    macro_rules! path_tests {
        ($($name:ident: $value:expr,)*) => {
        $(
            #[test]
            fn $name() {
                let (to_test, path, expected) = $value;
                let p = Path::new(path);
                assert_eq!(expected, path_match(to_test, p));
            }
        )*
        }
    }

    path_tests! {
        path_1: ("test", "test.so", true),
        path_2: ("pam", "test.so", false),
        path_3: ("pam", "libpam.so", true),
        path_4: ("krb5", "libkrb5.so.3.3", true),
        path_5: ("pthread-2.17", "/usr/bin/libpthread-2.17.so", true),
        path_6: ("krb5", "/usr/lib64/libkrb5.so.3.3", true),
    }
}
